from dataclasses import dataclass

import torch

from .utils import gather, scatter_set


@dataclass(kw_only=True)
class SynergyRouterOutput:
    tokens_logit: torch.Tensor
    tokens_sigma: torch.Tensor
    tokens_bool: torch.Tensor
    middle_tokens_index: torch.Tensor
    middle_tokens_logit: torch.Tensor
    middle_tokens_sigma: torch.Tensor


class SynergyRouter(torch.nn.Module):
    def __init__(self, *,
        m_size: int | None = None,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ):
        super().__init__()
        if m_size is None:
            self.linear = torch.nn.LazyLinear(1, dtype=dtype, device=device)
        else:
            self.linear = torch.nn.Linear(m_size, 1, dtype=dtype, device=device)

    def forward(self, *,
        tokens_in_emb: torch.Tensor,
        tokens_out_mask: torch.Tensor | None = None,
        middle_tokens_n: int,
        temperature: torch.Tensor | float | None = None,
    ) -> SynergyRouterOutput:
        tokens_logit = self.linear(tokens_in_emb).squeeze(-1)
        tokens_logit = torch.where(tokens_out_mask, tokens_logit, -torch.inf) \
            if tokens_out_mask is not None else tokens_logit
        tokens_sigma = torch.sigmoid(tokens_logit)
        # [..., chunk_tokens_n]

        if (isinstance(temperature, (float, int)) and temperature == 0) or temperature is None:
            _, middle_tokens_index = torch.topk(tokens_logit, middle_tokens_n, dim=-1, sorted=False)
        elif isinstance(temperature, float) and temperature == float('inf'):
            tokens_weight = torch.ones_like(tokens_logit)  # sample uniformly
            middle_tokens_index = torch.multinomial(tokens_weight, middle_tokens_n, replacement=False)
        else:
            temperature = torch.asarray(temperature).to(tokens_logit)
            tokens_logit_clamped = torch.clamp(tokens_logit, min=-1000, max=1000)
            tokens_weight = torch.softmax(tokens_logit_clamped / temperature, dim=-1)
            middle_tokens_index = torch.multinomial(tokens_weight, middle_tokens_n, replacement=False)
        middle_tokens_index, _ = torch.sort(middle_tokens_index, dim=-1)
        # [..., middle_tokens_n], int64

        middle_tokens_bool = torch.ones_like(middle_tokens_index, dtype=torch.bool)
        tokens_bool = torch.zeros_like(tokens_sigma, dtype=torch.bool)
        tokens_bool = scatter_set(tokens_bool, middle_tokens_bool, middle_tokens_index, dim=-1)
        # [..., middle_tokens_n], bool

        middle_tokens_logit = gather(tokens_logit, middle_tokens_index, dim=-1)
        middle_tokens_sigma = gather(tokens_sigma, middle_tokens_index, dim=-1)
        # [..., middle_tokens_n]

        return SynergyRouterOutput(
            tokens_logit=tokens_logit,
            tokens_sigma=tokens_sigma,
            tokens_bool=tokens_bool,
            middle_tokens_index=middle_tokens_index,
            middle_tokens_logit=middle_tokens_logit,
            middle_tokens_sigma=middle_tokens_sigma)
